{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d006b2ea-9dfe-49c7-88a9-a5a0775185fd",
   "metadata": {},
   "source": [
    "# Additional End of week Exercise - week 2\n",
    "\n",
    "Now use everything you've learned from Week 2 to build a full prototype for the technical question/answerer you built in Week 1 Exercise.\n",
    "\n",
    "This should include a Gradio UI, streaming, use of the system prompt to add expertise, and the ability to switch between models. Bonus points if you can demonstrate use of a tool!\n",
    "\n",
    "If you feel bold, see if you can add audio input so you can talk to it, and have it respond with audio. ChatGPT or Claude can help you, or email me if you have questions.\n",
    "\n",
    "I will publish a full solution here soon - unless someone beats me to it...\n",
    "\n",
    "There are so many commercial applications for this, from a language tutor, to a company onboarding solution, to a companion AI to a course (like this one!) I can't wait to see your results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b3acf7",
   "metadata": {},
   "source": [
    "# NOTE: Tool call to course notebooks \n",
    "\n",
    "This ended up being a bit more complex than I expected, so I only impleneted tool calling for chatgpt (not claude and gemini) as I had planned\n",
    "\n",
    "I ran into some problems getting streaming to work with tool calling. \n",
    "\n",
    "Also, the current implementation is not pretty :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07e7793-b8f5-44f4-aded-5562f633271a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# base imports\n",
    "\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "import anthropic\n",
    "import gradio as gr\n",
    "\n",
    "load_dotenv(override=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7f02f5c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantaite clients and set system prompt\n",
    "\n",
    "openai = OpenAI()\n",
    "claude = anthropic.Anthropic()\n",
    "\n",
    "SYSTEM_PROMPT = \"\\n\".join([\n",
    "    \"You are a helpful technical tutor who answers questions about python code, software engineering, data science and LLMs\",\n",
    "    \"You have access to a notebook_search tool that can search the course notebooks for relevant information to the user's question\",\n",
    "    \"You always keep your answers concise and to the point\",\n",
    "])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f92b2fa5",
   "metadata": {},
   "source": [
    "## This is the tool\n",
    "An index of embeddings for the course material - in this case just Week 2. But we could expand it to cover all the course material, so we can ask questions about it, and find references to things we forgot :)\n",
    "\n",
    "We can provide the URL to the Notebooks class that we want to query access to\n",
    "\n",
    "We opt out of the community contributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f55cbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import faiss\n",
    "\n",
    "# Set path to course notebooks\n",
    "NOTEBOOK_DIR = Path('~/code/llm_engineering/week2').expanduser()\n",
    "\n",
    "# Set embedding model (we could also use openai's embedding model)\n",
    "EMBED_MODEL = \"all-MiniLM-L6-v2\"\n",
    "\n",
    "\n",
    "class Notebooks:\n",
    "    def __init__(self, notebook_dir: Path = None):\n",
    "        self.embed_model = SentenceTransformer(EMBED_MODEL)\n",
    "        if notebook_dir:\n",
    "            self.load_notebooks(notebook_dir)\n",
    "\n",
    "    # Load all notebooks to memory\n",
    "    def load_notebooks(self, notebook_dir: Path):\n",
    "        print('Reading from', notebook_dir)\n",
    "        self.docs = []\n",
    "        for notebook_path in notebook_dir.rglob(\"*.ipynb\"):\n",
    "            if 'community-contributions' in str(notebook_path):\n",
    "                continue\n",
    "            print(notebook_path)\n",
    "\n",
    "            data = json.loads(notebook_path.read_text())\n",
    "            \n",
    "            # Include both markdown and code if available\n",
    "            cells = []\n",
    "            for cell in data.get(\"cells\", []):\n",
    "                if cell.get(\"cell_type\") == \"markdown\":\n",
    "                    cells.append(\"\".join(cell[\"source\"]))\n",
    "                elif cell.get(\"cell_type\") == \"code\":\n",
    "                    code = \"\".join(cell[\"source\"])\n",
    "                    cells.append(f\"```python\\n{code}\\n```\")\n",
    "                    if \"outputs\" in cell:\n",
    "                        for output in cell[\"outputs\"]:\n",
    "                            if \"text\" in output:\n",
    "                                cells.append(\"\".join(output[\"text\"]))\n",
    "            \n",
    "            text = \"\\n\\n\".join(cells).strip()\n",
    "            \n",
    "            if text:\n",
    "                self.docs.append({\n",
    "                    \"path\": str(notebook_path.relative_to(notebook_dir)),\n",
    "                    \"text\": text\n",
    "                })\n",
    "        \n",
    "        self._build_notebook_retriever()\n",
    "\n",
    "    # Build FAISS index for retreival\n",
    "    def _build_notebook_retriever(self):\n",
    "        print('Building search index')\n",
    "        texts = [d[\"text\"] for d in self.docs]\n",
    "\n",
    "        # Transform notebook text into embeddings\n",
    "        embeddings = self.embed_model.encode(texts, convert_to_numpy=True, show_progress_bar=True)\n",
    "\n",
    "        self.doc_index = faiss.IndexFlatL2(embeddings.shape[1])\n",
    "        self.doc_index.add(embeddings)\n",
    "\n",
    "    # Returns top n most similar notebook-markdown snippets\n",
    "    def search(self, query: str, top_n: int = 3, max_distance: float = None):\n",
    "        print('Looking for', query)\n",
    "        # compute embeddings for the query\n",
    "        embeddings = self.embed_model.encode([query], convert_to_numpy=True)\n",
    "        \n",
    "        # search the index\n",
    "        distances, indices = self.doc_index.search(embeddings, top_n)\n",
    "\n",
    "        # compile results\n",
    "        results = []\n",
    "        for dist, idx in zip(distances[0], indices[0]):\n",
    "            if max_distance is not None and dist > max_distance:\n",
    "                continue\n",
    "            \n",
    "            doc = self.docs[idx]\n",
    "            excerpt = doc[\"text\"]\n",
    "            if len(excerpt) > 500:\n",
    "                excerpt = excerpt[:500].rsplit(\"\\n\", 1)[0] + \"…\"\n",
    "            \n",
    "            results.append({\n",
    "                \"source\": doc[\"path\"],\n",
    "                \"excerpt\": excerpt,\n",
    "                \"score\": float(dist) # lower socre is more similar in L2 space\n",
    "            })\n",
    "        \n",
    "        return results\n",
    "    \n",
    "    def as_tool(self):\n",
    "        return { \n",
    "            \"type\": \"function\", \n",
    "            \"function\": {\n",
    "                \"name\": \"notebook_search\",\n",
    "                \"description\": \"Searches the course notebooks and returns relevant excerpts with paths.\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"query\": {\n",
    "                            \"type\": \"string\", \n",
    "                            \"description\": \"What to look for in the course notebooks\"\n",
    "                        },\n",
    "                        \"top_n\": {\n",
    "                            \"type\":\"integer\",\n",
    "                            \"description\":\"How many course notebook passages to return\", \n",
    "                            \"default\": 3\n",
    "                        }\n",
    "                    },\n",
    "                    \"required\": [\"query\"],\n",
    "                    \"additionalProperties\": False\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "        \n",
    "    \n",
    "notebooks = Notebooks(NOTEBOOK_DIR)\n",
    "\n",
    "\n",
    "def notebook_search(query, top_n=3):\n",
    "    return notebooks.search(query, top_n)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce7608bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test tool here\n",
    "\n",
    "notebooks.search(\"Gradio\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a214dd2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODELS = dict(\n",
    "    gpt='gpt-4o-mini',\n",
    "    claude='claude-3-haiku-20240307',\n",
    ")\n",
    "\n",
    "def get_interactions(message, history):\n",
    "    messages = []\n",
    "    for user_msg, bot_msg in history:\n",
    "        messages.append({\"role\":\"user\", \"content\":user_msg})\n",
    "        messages.append({\"role\":\"assistant\", \"content\":bot_msg})\n",
    "    messages.append({\"role\":\"user\", \"content\":message})\n",
    "    return messages\n",
    "\n",
    "\n",
    "def get_chatgpt_stream(model, message, history):\n",
    "    print(f\"Getting OpenAi stream, using {model}\")\n",
    "    interactions = get_interactions(message, history)\n",
    "    messages = [{\"role\": \"system\", \"content\": SYSTEM_PROMPT}] + interactions\n",
    "\n",
    "    stream = openai.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=messages,\n",
    "        temperature=0.5,\n",
    "        stream=True,\n",
    "        tools=[\n",
    "            notebooks.as_tool()\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    tool_call = None\n",
    "    fn_name = None\n",
    "    fn_args = \"\"\n",
    "    tool_call_id = None\n",
    "    buffer = \"\"\n",
    "    \n",
    "    for chunk in stream:\n",
    "        delta = chunk.choices[0].delta\n",
    "\n",
    "        # Handle normal content\n",
    "        if delta and delta.content:\n",
    "            buffer += delta.content or ''\n",
    "            yield buffer\n",
    "\n",
    "        # Handle tool call\n",
    "        if delta and delta.tool_calls:\n",
    "            tool_call = delta.tool_calls[0]\n",
    "            if tool_call.id:\n",
    "                tool_call_id = tool_call.id\n",
    "            if tool_call.function.name:\n",
    "                fn_name = tool_call.function.name\n",
    "            if tool_call.function.arguments:\n",
    "                fn_args += tool_call.function.arguments\n",
    "            yield buffer  # Yield to keep Gradio updated\n",
    "    \n",
    "    if fn_name == \"notebook_search\" and fn_args and tool_call_id:\n",
    "        print('Tool call to ', fn_name)\n",
    "\n",
    "        args = json.loads(fn_args)\n",
    "        result = notebook_search(**args)  # Returns list of dicts\n",
    "        result_str = json.dumps(result, indent=2)\n",
    "        print(\"Tool result:\", result_str)\n",
    "\n",
    "        # Append assistant message with tool call\n",
    "        messages.append({\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": None,\n",
    "            \"tool_calls\": [\n",
    "                {\n",
    "                    \"id\": tool_call_id,\n",
    "                    \"type\": \"function\",\n",
    "                    \"function\": {\n",
    "                        \"name\": fn_name,\n",
    "                        \"arguments\": fn_args\n",
    "                    }\n",
    "                }\n",
    "            ]\n",
    "        })\n",
    "\n",
    "        messages.append({\n",
    "            \"role\": \"tool\",\n",
    "            \"content\": result_str,\n",
    "            \"tool_call_id\": tool_call_id\n",
    "        })\n",
    "        messages.append({\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": \"Make sure you reference the source notebook in your answer.\",\n",
    "        })\n",
    "\n",
    "        # Follow-up chat call\n",
    "        followup_stream = openai.chat.completions.create(\n",
    "            model=model,\n",
    "            messages=messages,\n",
    "            temperature=0.5,\n",
    "            stream=True\n",
    "        )\n",
    "\n",
    "        # Stream follow-up response\n",
    "        for chunk in followup_stream:\n",
    "            delta = chunk.choices[0].delta\n",
    "            if delta.content:\n",
    "                buffer += delta.content or \"\"\n",
    "                yield buffer\n",
    "\n",
    "\n",
    "def get_claude_stream(model, message, history):\n",
    "    print(f\"Getting Claude stream, using {model}\")\n",
    "    interactions = get_interactions(message, history)\n",
    "\n",
    "    with claude.messages.stream(\n",
    "        model=model,\n",
    "        messages=interactions,\n",
    "        max_tokens=500,\n",
    "        system=SYSTEM_PROMPT,\n",
    "    ) as stream:\n",
    "        buffer = \"\"\n",
    "        for delta in stream.text_stream:\n",
    "            buffer += delta\n",
    "            yield buffer\n",
    "\n",
    "\n",
    "def chat(model_selector, message, history):\n",
    "    model = MODELS.get(model_selector)\n",
    "    if not model:\n",
    "        raise ValueError(f\"Invalid model: {model_selector}\")\n",
    "    \n",
    "    reply = \"\"\n",
    "    if model_selector == 'gpt':\n",
    "        for partial in get_chatgpt_stream(model, message, history):\n",
    "            reply = partial\n",
    "            yield history + [(message, reply)]\n",
    "\n",
    "    elif model_selector == 'claude':\n",
    "        for partial in get_claude_stream(model, message, history):\n",
    "            reply = partial\n",
    "            yield history + [(message, reply)]\n",
    "    \n",
    "\n",
    "with gr.Blocks() as demo:\n",
    "    model_selector = gr.Dropdown(\n",
    "        choices=MODELS.keys(),\n",
    "        value=\"gpt\", \n",
    "        label=\"Pick Model\",\n",
    "    )\n",
    "    chatbot = gr.Chatbot()\n",
    "    txt = gr.Textbox(placeholder=\"Ask about python\", show_label=False)\n",
    "    txt.submit(\n",
    "        fn=chat,\n",
    "        inputs=[model_selector, txt, chatbot],\n",
    "        outputs=[chatbot],\n",
    "    ).then(\n",
    "        fn=lambda: \"\",\n",
    "        inputs=None,\n",
    "        outputs=txt\n",
    "    )\n",
    "\n",
    "    clear = gr.Button(\"Clear\")\n",
    "    clear.click(lambda: None, None, chatbot, queue=False)\n",
    "\n",
    "demo.launch()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc128d47",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llms",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
